'''
Time Surface Generation from RAW files
Simplified version - with viridis color mode
'''
import cv2
import numpy as np
import os
import argparse
from tqdm import tqdm
from metavision_core.event_io import EventsIterator

np.random.seed(0)

def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--raw_file_dir', type=str,
                        default='H:\Datasets/val/normal',
                        help='The path of a raw dir')
    parser.add_argument('--save_dir', type=str,
                        default='H:\Datasets/val/timesurface',
                        help='Output directory for time surface images')
    parser.add_argument('--time_interval_us', type=int, default=50000,
                        help='Time interval in microseconds')
    return parser

def get_reader(file_path, time_interval_us):
    assert os.path.exists(file_path), 'The file \'{}\' does not exist'.format(file_path)
    mv_iterator = EventsIterator(input_path=file_path, delta_t=time_interval_us)
    return mv_iterator

def create_viridis_colormap(time_surface):
    """优化版本的viridis色彩映射"""
    viridis_colors = np.array([
        [68, 1, 84], [59, 82, 139], [33, 144, 140], 
        [93, 201, 99], [253, 231, 37]
    ]) / 255.0
    
    mask = time_surface > 0
    values = time_surface[mask]
    
    color_indices = values * (len(viridis_colors) - 1)
    idx_low = np.floor(color_indices).astype(int)
    idx_high = np.minimum(idx_low + 1, len(viridis_colors) - 1)
    alpha = color_indices - idx_low
    
    colors = ((1 - alpha[..., None]) * viridis_colors[idx_low] + 
              alpha[..., None] * viridis_colors[idx_high])
    
    colored_surface = np.zeros((time_surface.shape[0], time_surface.shape[1], 3))
    colored_surface[mask] = colors
    
    return (colored_surface * 255).astype(np.uint8)[:, :, ::-1]  # RGB to BGR
# def create_viridis_colormap(time_surface):
#     """超快版本 - 使用预计算的查找表"""
#     if not hasattr(create_viridis_colormap, 'lut'):
#         # 预计算256级查找表
#         viridis_colors = np.array([
#             [68, 1, 84], [59, 82, 139], [33, 144, 140], 
#             [93, 201, 99], [253, 231, 37]
#         ]) / 255.0
        
#         lut = np.zeros((256, 3))
#         for i in range(256):
#             value = i / 255.0
#             color_idx = value * (len(viridis_colors) - 1)
#             idx_low = int(np.floor(color_idx))
#             idx_high = min(idx_low + 1, len(viridis_colors) - 1)
#             alpha = color_idx - idx_low
#             lut[i] = (1 - alpha) * viridis_colors[idx_low] + alpha * viridis_colors[idx_high]
        
#         create_viridis_colormap.lut = lut
    
#     # 直接查表，超快速度
#     time_surface_uint8 = (time_surface * 255).astype(np.uint8)
#     colored = create_viridis_colormap.lut[time_surface_uint8]
    
#     return (colored * 255)[:, :, ::-1].astype(np.uint8)  # RGB to BGR

def create_time_surface(events, height, width):
    """
    Create time surface representation from events with viridis colormap
    
    Args:
        events: event data with 't', 'x', 'y', 'p' fields
        height, width: image dimensions
    
    Returns:
        time_surface: viridis colored time surface image
    """
    # Initialize time surface with zeros
    time_surface = np.zeros((height, width), dtype=np.float64)
    
    if len(events) == 0:
        return np.zeros((height, width, 3), dtype=np.uint8)
    
    # Filter valid events
    valid_mask = (events['x'] >= 0) & (events['x'] < width) & \
                 (events['y'] >= 0) & (events['y'] < height)
    valid_events = events[valid_mask]
    
    if len(valid_events) == 0:
        return np.zeros((height, width, 3), dtype=np.uint8)
    
    # For each pixel, keep only the latest event time
    for i in range(len(valid_events)):
        x, y, t = valid_events['x'][i], valid_events['y'][i], valid_events['t'][i]
        if time_surface[y, x] < t:
            time_surface[y, x] = t
    
    # Linear normalization: map time range to [0, 1]
    non_zero_mask = time_surface > 0
    if np.any(non_zero_mask):
        min_time = np.min(time_surface[non_zero_mask])
        max_time = np.max(time_surface[non_zero_mask])
        if max_time > min_time:
            time_surface[non_zero_mask] = (time_surface[non_zero_mask] - min_time) / (max_time - min_time)
        else:
            time_surface[non_zero_mask] = 1.0
    
    # Apply viridis colormap
    return create_viridis_colormap(time_surface)

if __name__ == '__main__':
    # Get params
    args, _ = parse_argument().parse_known_args(None)
    print(args)

    # Find all raw files
    raw_path_list = []
    for root, dirs, files in os.walk(args.raw_file_dir):
        for file in files:
            if file.endswith('.raw'):
                raw_path_list.append(os.path.join(root, file))
    
    print(f"Found {len(raw_path_list)} raw files")
    print("Using viridis color mode")

    for raw_path in tqdm(raw_path_list):
        assert os.path.exists(raw_path)
        base_name = os.path.basename(raw_path).split('.')[0]
        print(f"Processing: {base_name}")

        # Create output directory
        save_path = os.path.join(args.save_dir, f"{base_name}")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        reader = get_reader(raw_path, args.time_interval_us)
        height, width = reader.get_size()
        
        print(f"Image dimensions: {width}x{height}")

        frame_cnt = 0
        for evs in reader:
            if len(evs) == 0:
                continue
                
            min_t, max_t = np.min(evs['t']), np.max(evs['t'])
            
            # Create time surface representation with viridis colormap
            time_surface = create_time_surface(evs, height, width)
            
            # Save image
            filename = f'{save_path}/{frame_cnt:06d}_{min_t}_{max_t}.png'
            cv2.imwrite(filename, time_surface)
            
            frame_cnt += 1
            
            if frame_cnt % 100 == 0:
                print(f"Processed {frame_cnt} frames")
        
        print(f"Completed {base_name}: {frame_cnt} frames generated")

    print("Time surface generation completed!")